#include "MassProperties.h"
#include "RigidBodyNode.h"

#include "RigidBody.h"
#include "Mat6x6.h"
#include "Index.h"
#include "Force.h"
#include "MultibodySystem.h"
#include "FastOps.h"
#include "EOMData.h"
#include "ABAData.h"

#include "Joint.h"
#include "Ball.h"
#include "Pin.h"
#include "PrismaticZ.h"
#include "PrismaticX.h"
#include "XYZJoint.h"
#include "Cylinder.h"
#include "FreeJoint.h"
#include "Weld.h"

#include <cstdlib>
#include <string>

using namespace std;
using namespace RSIM;

RigidBody::RigidBody(const char *str, 
                const MassProperties& MP, 
                const Transform& Body_T_BF, 
                const Transform& Parent_T_PM):	Mass_(MP.Mass_),Inertia_(MP.Inertia_),J_(NULL),
                                                        T_BF_(Body_T_BF),T_PM_(Parent_T_PM),
                                                        T_FM_(NULL),R_NF_(),prF_C_F_(),
                                                        spF_(NULL),EOM_(NULL),ABA_(NULL),w_(0.0),v_(0.0),
                                                        rk_(0.0),sAkt_(0.0),
                                                        nChildren_(0),JT_(Invalid){
        if(!strcmp(str,"Ball")){
                JT_ = BallJoint;
        }
        else if(!strcmp(str,"Pin")){
                JT_ = PinJoint;
        }
        else if(!strcmp(str,"PrismaticZ")){
                JT_ = PrismaticZJoint;
        }
        else if(!strcmp(str,"PrismaticX")){
                JT_ = PrismaticXJoint;
        }
        else if(!strcmp(str,"Cylinder")){
                JT_ = CylinderJoint;
        }
        else if(!strcmp(str,"XYZ")){
                JT_ = XYZJoint;
        }
        else if(!strcmp(str,"Free")){
                JT_ = FreeJoint;
        }
        else if(!strcmp(str,"Weld")){
                JT_ = WeldJoint;
        }
        else{
                cout<<"\nInvalid option. Possible joint types are:\n <Ball>,\n <Pin>,\n "<<
                "<PrismaticZ>,\n <PrismaticX>,\n <Cylinder>,\n <XYZ>,\n <Free>,\n <Weld>,\n";
                exit(0);
        }
}

///////////////////////////////////////////////////////

RigidBody::RigidBody(const RigidBody& RB):Mass_(RB.getMass()),Inertia_(RB.getInertia()),
                J_(NULL),T_BF_(RB.getTransformBF()),
                T_PM_(RB.getTransformPM()),
                T_FM_(NULL),R_NF_(RB.getRotationNF()),		
                prF_C_F_(RB.getRotationFCF()),spF_(NULL),
                EOM_(NULL),ABA_(NULL),w_(0.0),rk_(0.0),sAkt_(0.0),nChildren_(RB.getNChildren()),
                JT_(RB.getJointType()){
}

///////////////////////////////////////////////////////

void RigidBody::initJoint(const int& NodeIndex){
        if(JT_== BallJoint)
                this->J_ = Ball::New(NodeIndex);
        else if(JT_==PinJoint){
                this->J_ = Pin::New(NodeIndex);
        }
        else if(JT_==PrismaticZJoint){
                this->J_=PrismaticZ::New(NodeIndex);
        }
        else if(JT_==PrismaticXJoint){
                this->J_=PrismaticX::New(NodeIndex);
        }
        else if(JT_==CylinderJoint){
                this->J_=Cylinder::New(NodeIndex);
        }
        else if(JT_==XYZJoint){
                this->J_=XYZ::New(NodeIndex);
        }
        else if(JT_==FreeJoint){
                this->J_=Free::New(NodeIndex);
        }
        else if(JT_==WeldJoint){
                this->J_=Weld::New(NodeIndex);
        }
        else{
                cout<<"\nInvalid Joint type\n";
                exit(0);
        }
}

///////////////////////////////////////////////////////

void RigidBody::setBodyRK(const double *rk){
        double *data = this->rk_.wPtr();
        data[0] = rk[0]; 
        data[1] = rk[1];
        data[2] = rk[2];
}

///////////////////////////////////////////////////////

void RigidBody::formEOM(const SystemAcceleration *SysAccl){
        
        ///////////////////////////////////////////////////////
        // Calculation of the known external forces acting on 
        // this body and the gyroscopic terms. 
        ///////////////////////////////////////////////////////

        double *fptr = this->EOM_->wgetdFk1()->wPtr();
        // Initialize to zero
        fptr[0]=0.0;fptr[1]=0.0;fptr[2]=0.0;
        fptr[3]=0.0;fptr[4]=0.0;fptr[5]=0.0;
        
        // External spatial force if applicable
        if(this->spF_){
                this->spF_->applySpatialForce(fptr);
        }
        
        // System acceleration if applicable
        if(SysAccl){
                double grav_force[3];
                double grav_torque[3];
                
                FastOps::transpose_mult_mat33_vec3(this->R_NF_.Ptr(), SysAccl->getGravity().Ptr(), grav_force);
                grav_force[0] *= this->Mass_;
                grav_force[1] *= this->Mass_;
                grav_force[2] *= this->Mass_;
                
                FastOps::cross(this->rk_.Ptr(),grav_force,grav_torque);
                
                fptr[0] += grav_torque[0];
                fptr[1] += grav_torque[1];
                fptr[2] += grav_torque[2];
                
                fptr[3] += grav_force[0];
                fptr[4] += grav_force[1];
                fptr[5] += grav_force[2];
        }	
        ///////////////////////////////////////////////////////	
        // Computing {[wx]J.w, m(cross(w,cross(w,r)))}^T
        ///////////////////////////////////////////////////////	
        
        double top[3];
        this->Wx_J_W(this->w_.Ptr(),this->EOM_->getJk1()->Ptr(),top);
        
        double bot[3];
        FastOps::wxwxr(this->w_.Ptr(),this->rk_.Ptr(),bot);
                
        fptr[0] -= top[0];
        fptr[1] -= top[1];
        fptr[2] -= top[2];
        
        fptr[3] -= this->Mass_*bot[0];
        fptr[4] -= this->Mass_*bot[1];
        fptr[5] -= this->Mass_*bot[2];
        
        double *fkc1_ptr = this->EOM_->wgetFkc1()->wPtr();
        fkc1_ptr[0]=0.0;fkc1_ptr[1]=0.0;fkc1_ptr[2]=0.0;
        fkc1_ptr[3]=0.0;fkc1_ptr[4]=0.0;fkc1_ptr[5]=0.0;		
        
        if(ABA_){
                ABA_->wgetIk3()[0].copy(EOM_->getdMk1()->Ptr());
                ABA_->wgetFk3()[0].copy(EOM_->getdFk1()->Ptr());
        }
}

///////////////////////////////////////////////////////

/** P = [Wx][C][W] */
void RigidBody::Wx_J_W(const double *w, const double *c, double *p){
        p[0]= -w[0]*(c[3]*w[2] - c[6]*w[1]) - w[1]*(c[4]*w[2] - c[7]*w[1]) - w[2]*(c[5]*w[2] - c[8]*w[1]);
        p[1]=  w[0]*(c[0]*w[2] - c[6]*w[0]) + w[1]*(c[1]*w[2] - c[7]*w[0]) + w[2]*(c[2]*w[2] - c[8]*w[0]);
        p[2]= -w[0]*(c[0]*w[1] - c[3]*w[0]) - w[1]*(c[1]*w[1] - c[4]*w[0]) - w[2]*(c[2]*w[1] - c[5]*w[0]);
}

///////////////////////////////////////////////////////

void RigidBody::sweepForwardABA(const RigidBodyNode *Parent,const int& ChildID, State& state){
        
        this->J_->sweepForward(Parent, ChildID,this->sAkt_,this->prF_C_F_,this->ABA_,this->betak_,this->EOM_,state);
}

///////////////////////////////////////////////////////

void RigidBody::sweepBackwardRNEA(const int& ChildID, RigidBodyNode *Parent, State& state){
                
        double *body_fkc1 = this->EOM_->wgetFkc1()->wPtr();
        const double *dfk1_ptr = this->EOM_->getdFk1()->Ptr();
        
        double M_x_A[6];
        FastOps::mult_mat66_vec6(EOM_->getdMk1()->Ptr(), EOM_->getAk1()->Ptr(), M_x_A);
        
        // This completes the computation of constraint force acting on this body
        body_fkc1[0] += M_x_A[0]-dfk1_ptr[0];
        body_fkc1[1] += M_x_A[1]-dfk1_ptr[1];
        body_fkc1[2] += M_x_A[2]-dfk1_ptr[2];
        body_fkc1[3] += M_x_A[3]-dfk1_ptr[3];
        body_fkc1[4] += M_x_A[4]-dfk1_ptr[4];
        body_fkc1[5] += M_x_A[5]-dfk1_ptr[5];
        
        if(Parent){
// 		const Transform& T_FM = Parent->data_->getBodyTransformFM()[ChildID];
                double *parent_fkc1 = Parent->data_->wgetEOMData()->wgetFkc1()->wPtr();
                                
                double sFkc1[12];
                FastOps::spatial_basis_shift_vec(this->prF_C_F_.Ptr(),body_fkc1,sFkc1);
                FastOps::spatial_shift_vec6(sFkc1, this->betak_, sFkc1+6);
                
                parent_fkc1[0] += sFkc1[6];
                parent_fkc1[1] += sFkc1[7];
                parent_fkc1[2] += sFkc1[8];
                parent_fkc1[3] += sFkc1[9];
                parent_fkc1[4] += sFkc1[10];
                parent_fkc1[5] += sFkc1[11];
        }
}

///////////////////////////////////////////////////////

void RigidBody::sweepForwardRNEA(const RigidBodyNode *Parent, 				      
                                const int& ChildID,
                                const SystemAcceleration *SysAccl,
                                RigidBodyNode *Body,
                                State& state){

        // Form newton euler equations of motion
        this->formEOM(SysAccl);
        
        // Calculate akt. this includes the udot terms since motion is always
        // prescribed in RNEA.
        this->J_->calcAKT(Parent, Body, ChildID, this->betak_, state, this->sAkt_);
        
        const double *sAkt_ptr = sAkt_.Ptr();
        
        double *ak1_ptr = this->EOM_->wgetAk1()->wPtr(); 
        ak1_ptr[0] = 0.0;ak1_ptr[1] = 0.0;ak1_ptr[2] = 0.0;
        ak1_ptr[3] = 0.0;ak1_ptr[4] = 0.0;ak1_ptr[5] = 0.0;
        
        if(Parent){
                double tmp[6];
                FastOps::transpose_skbeta_x_V(Parent->data_->getEOMData()->getAk1()->Ptr(), this->betak_, tmp);
                FastOps::transpose_spatial_basis_shift_vec(this->prF_C_F_.Ptr(),tmp,ak1_ptr);
                
                ak1_ptr[0] += sAkt_ptr[0];	ak1_ptr[1] += sAkt_ptr[1];	ak1_ptr[2] += sAkt_ptr[2];
                ak1_ptr[3] += sAkt_ptr[3];	ak1_ptr[4] += sAkt_ptr[4];	ak1_ptr[5] += sAkt_ptr[5];
        }
        else{
                J_->calc_H_x_UDot(state, ak1_ptr);		
        }
}

///////////////////////////////////////////////////////

RigidBody::~RigidBody(){
        delete [] T_FM_;
        delete spF_;
        delete ABA_;
        delete EOM_;
        delete J_;
}

///////////////////////////////////////////////////////

void RigidBody::getConstraintForce(double *Fc) const{
        const double *ptr = this->EOM_->getFkc1()->Ptr();
        Fc[0] = ptr[0];	Fc[1] = ptr[1];	Fc[2] = ptr[2];
        Fc[3] = ptr[3];	Fc[4] = ptr[4];	Fc[5] = ptr[5];
}

///////////////////////////////////////////////////////

void RigidBody::setRotationNF(){
        // R_NF = R_PM * R_MF
        // Frame P in this case is the newtonian frame
        FastOps::mult_mat33_mat33(this->T_PM_.R().Ptr(), this->J_->getRotationMF().Ptr(), this->R_NF_.wPtr());
}

///////////////////////////////////////////////////////

Vect3 RigidBody::expressAngularVelocityInGround() const{
        return Vect3(R_NF_*w_);
}

///////////////////////////////////////////////////////

void RigidBody::initEOM(){
        
        // Allocate space and compute the constant 
        // terms in the Newton-Euler equations of motion
        // of this body
        this->EOM_ = new EOMData;
        this->EOM_->initEOM(Mass_,Inertia_,T_BF_);
}

///////////////////////////////////////////////////////

void RigidBody::initForwardDynamics(const int& nchildren){
        #ifdef RSIM_SAFE
        if(!this->EOM_){
                cout<<"\n MultibodySystem::initEOM() must be called before calling initForwardDynamics()\n";
                exit(0);
        }
        #endif
        this->ABA_ = new ABAData;
        this->ABA_->initForwardDynamics(this->J_);
}

///////////////////////////////////////////////////////

void RigidBody::printJointType(){
        J_->printJointType();
}

///////////////////////////////////////////////////////

void RigidBody::printInfo(){
        cout<<"Mass="<<Mass_;
}

